import matplotlib.pyplot as plt
import numpy as np
from sklearn.tree import DecisionTreeRegressor

import getData
from TwoStageTrAdaBoostR2 import TwoStageTrAdaBoostR2


class Tradaboost_r2():

    def __init__(self):
        self.n_estimators = 100
        self.steps = 50
        self.fold = 5
        self.random_state = np.random.RandomState(42)
        self.data = getData.GetData()

    def fit(self, data):
        data.getData()
        x_source = data.source_x_train
        y_source = data.source_y_train
        y_source1 = y_source[:, 0].reshape(-1)
        y_source2 = y_source[:, 1].reshape(-1)

        x_target_train = data.target_x_train
        y_target_train = data.target_y_train
        y_target_train1 = y_target_train[:, 0].reshape(-1)
        y_target_train2 = y_target_train[:, 1].reshape(-1)


        X = np.concatenate((x_source, x_target_train))
        y = np.concatenate((y_source1, y_target_train1))
        sample_size = [x_source.shape[0], x_target_train.shape[0]]
        self.model1 = TwoStageTrAdaBoostR2(base_estimator=DecisionTreeRegressor(max_depth=6),
                                          n_estimators=self.n_estimators,
                                          sample_size=sample_size,
                                          steps=self.steps,
                                          fold=self.fold,
                                          learning_rate=0.1,
                                          random_state=self.random_state)
        self.model1.fit(X, y)

        # ===============================================================
        X = np.concatenate((x_source, x_target_train))
        y = np.concatenate((y_source2, y_target_train2))
        sample_size = [x_source.shape[0], x_target_train.shape[0]]
        self.model2 = TwoStageTrAdaBoostR2(base_estimator=DecisionTreeRegressor(max_depth=6),
                                           n_estimators=self.n_estimators,
                                           sample_size=sample_size,
                                           steps=self.steps,
                                           fold=self.fold,
                                           learning_rate=0.1,
                                           random_state=self.random_state)
        self.model2.fit(X, y)



    def predict(self, x_test):
        y_pre1 = self.model1.predict(x_test)[:, np.newaxis]
        y_pre2 = self.model2.predict(x_test)[:, np.newaxis]
        y_pre = np.hstack((y_pre1, y_pre2))
        print(y_pre)
        return y_pre


if __name__ == '__main__':
    tra = Tradaboost_r2()
    data = getData.GetData()
    tra.fit(data)

    x_test = data.target_x_test
    y_pre = tra.predict(x_test)
    y_test = data.target_y_test
    for i in range(y_pre.shape[1]):
        plt.figure()
        plt.plot(x_test, y_pre[:, i], color="red")
        plt.plot(x_test, y_test[:, i])
        plt.show()
